import torch
import numpy as np
import pandas as pd

df = pd.read_csv('/home/dl/datasets/celeba/Anno/list_attr_celeba.txt', sep=r'\s+', skiprows=1)
df.index.name='image_name'

df.replace(-1, 0, inplace=True)

status_df = pd.read_csv('/home/dl/datasets/celeba/Eval/list_eval_partition.txt', sep=r'\s+', names=['image_name', 'status'])
status_df = status_df.set_index('image_name')
id2status = {0:'train', 1:'val', 2:'test'}
status_df.status = status_df.status.map(lambda x: id2status[x])

info_df = df.join(status_df)
info_df.to_csv('./data/info.csv')